#!/usr/bin/env bash

usage() { 
    cat <<EOF
Usage: $0 [-u <20|22>] [-l <cn|global>] [-c] [-f]
        -u ## for ubuntu version
        -l xx for language(cn for Chinese, global for English)
        -c    for cuda
        -f    for always(force) rebuild images

    default                      : $0 -u 22  -l global 
    for Chinese                  : $0 -l cn   ( : $0 -u 22  -l cn )
    for ubuntu 20.04             : $0 -u 20.04
    for cuda(+pytorch)           : $0 -c
    for Chinese + cuda(+pytorch) : $0 -l cn -c
EOF

    exit 1
}

u=22
l=global
force=""
cuda=""

# getopts "u:l:cfh:" # kvpair x:  ; only-key x
while getopts "u:l:cfh" OPTKEY; do
    echo "key=${OPTKEY} , arg=$OPTARG"
    case "${OPTKEY}" in
        u)
            u=${OPTARG}
            ((u == 20 || u == 22)) || usage
            ;;
        l)
            l=${OPTARG}
            [ "$l" == "cn" -o "$l" == "global" ] || usage
            ;;
        c)
            cuda=true
            ;;
        f)
            force=true
            ;;
        :)
            echo "o=$o, optarg=${OPTARG}"
            if [ "${OPTARG}" == "f" ];then
                force=true
            elif [ "${OPTARG}" == "c" ];then
                cuda=true
            fi
            ;;
        h)
            usage
            ;;
    esac
done

shift $((OPTIND-1))


cn="$l"
ubuntu="$u.04"

echo "$@"
echo "[I] run as : $0 -l $cn -u $u  ${cuda:+-c} ${force:+-f}"
read -p "ENTER to entinue:" input
if [ -n "$input" ];then
    usage

    exit 
fi

## getopt

# dim
## ubuntu : (default) 22.04 ,  20.04
## lang:  cn, global


is_image_exist(){
    if [ -z "$(docker images -q $1 2> /dev/null)" ]; then
        return 1
    fi
    return 0
}


build_run_cuda(){
    local tag rdp pytorch_rdp
    tag="${cn}"

    rdp="rdp_cuda:$tag"
    pytorch_rdp="rdp_pytorch:$tag"

    if [ "$force" == "true" ] || ! is_image_exist "$pytorch_rdp"; then
        echo "[1] build base : $rdp"
        if [ "$force" == "true" ] || ! is_image_exist "$rdp"; then
            (
                cd "src"
                docker build -f Dockerfile.$ubuntu.$tag -t "$rdp"  --build-arg IMAGE=nvidia/cuda --build-arg TAG=12.4.1-runtime-ubuntu$ubuntu .
            )
            if [ "$?" -ne 0 ];then
                exit 3
            fi
            echo
            echo
        else
            echo "  already exists"
            echo
            echo
        fi
        # build cuda pytorch
        echo '[2] build cuda pytorch'
        # cuda_rdp_based_python="rdp_cuda_python"
        (
            cd  "src/cuda_pytorch"
            bash  build_stage1.sh "$rdp" "$pytorch_rdp" "$tag"
        )
        if [ "$?" -ne 0 ];then
            exit 4
        fi
    fi

    ###
    # run
    bash ./_start_pytorch "$ubuntu" "$pytorch_rdp" force

}

build_run_without_cuda(){

    local tag rdp pytorch_rdp
    tag="${cn}"

    rdp="rdp_ubuntu:$tag"


    echo "[1] build base : $rdp"
    if [ "$force" == "true" ] || ! is_image_exist "$rdp"; then
        (
            cd "src"
            docker build -f Dockerfile.$ubuntu.$tag -t "$rdp"  --build-arg IMAGE=ubuntu --build-arg TAG=$ubuntu .
        )
        if [ "$?" -ne 0 ];then
            exit 3
        fi
        echo
        echo
    else
        echo "  already exists"
        echo
        echo
    fi


    ###
    # run
    bash ./_start "$ubuntu" "$rdp" force
}


(
    cd "$(dirname "$0")"
    set -x

    if [ -z "$cuda" ];then
        build_run_without_cuda
    else
        build_run_cuda
    fi
    
)

